#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <mpi.h>
#include <png.h>
#include<omp.h>
#define MASTER_PROCESS 0

extern int * initRandomWalk( int size, int blockNum, int threadsNum, int rank );

inline void setRGB (png_byte *ptr, float val);
int writeImage(char* filename, int *buffer, char* title, int Size);

/*
void print_malha( int* matriz, int size )
{
    int i, j;
    for ( i = 0; i < size; i++ ){
        for ( j = 0; j < size; j++ ){
                printf("%5d", matriz[j + i * size] );

        }
        printf("\n");
    }
}
*/

int main(int argc, char *argv[])
{
    // ./<executable_name> <tamanho_da_matriz> <número_de_particulas> <nome_do_arquivo_de_saída>
    if ( argc < 4 ){
        printf("Entre com os argumentos necessários.\n");
        printf("./<executable_name> <tamanho_da_matriz> <número_de_blocos> <numero_de_threads_por_bloco> <nome_do_arquivo_de_saída>\n");
        return EXIT_FAILURE;
    }
    
    
    int size = atoi(argv[1]); // recupera o tamanho da matriz
    int blockNum = atoi(argv[2]); // recupera o número de blocos
    int threadsNum = atoi(argv[3]); // recupera o número de threads por bloco
    char outputName[50] = "arquivo_saida.png";   /* Nome de saida padrão para a imagem a gerada */
    strcpy(outputName, argv[4]); // captura o nome do arquivo de saída
    int totalParticulas = blockNum * threadsNum; // o total de particulas é equivalente ao total de threads disparados

    int * randowWalk = NULL;
    int i, myrank, numprocs;
    int image;
    int * result;
    
    MPI_Init(&argc,&argv);
    MPI_Comm_size(MPI_COMM_WORLD,&numprocs);
    MPI_Comm_rank(MPI_COMM_WORLD, &myrank);
    
    // Inicializa RandomWalk
    double start = omp_get_wtime();
    randowWalk = initRandomWalk(size, blockNum, threadsNum/numprocs, myrank);
    
    /* Apenas o processo 0 (MASTER) aloca espaço para o resultado da redução */
    if ( myrank == MASTER_PROCESS )
        result = (int*) calloc ( size * size, sizeof(int));
    
    
    /* Redução de soma para o processo de rank 0 (MASTER_PROCESS) */
    MPI_Reduce( randowWalk, result, size*size, MPI_INT, MPI_SUM, MASTER_PROCESS, MPI_COMM_WORLD );
   
    double end = omp_get_wtime();
    printf("Total: %lf\n", end - start );
    
    if ( myrank == MASTER_PROCESS ){
        image = writeImage( outputName, result, "This is a Random Walk result's image", size);
        // printf("\n");
        // print_malha( result, size );
    }
    MPI_Finalize();
    
    free(randowWalk);
    return 0;
}

int writeImage(char* filename, int *buffer, char* title, int Size)
{
    int x, y;
    int maior = buffer[0];
    int menor = buffer[0];
    int code = 0;
    FILE *fp;
    png_structp png_ptr;
    png_infop info_ptr;
    png_bytep row;

    // Open file for writing (binary mode)
    fp = fopen(filename, "wb");
    if (fp == NULL) {
	    fprintf(stderr, "Could not open file %s for writing\n", filename);
	    code = 1;
	    goto finalise;
    }

    // Initialize write structure
    png_ptr = png_create_write_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);
    if (png_ptr == NULL) {
	    fprintf(stderr, "Could not allocate write struct\n");
	    code = 1;
	    goto finalise;
    }

    // Initialize info structure
    info_ptr = png_create_info_struct(png_ptr);
    if (info_ptr == NULL) {
	    fprintf(stderr, "Could not allocate info struct\n");
	    code = 1;
	    goto finalise;
    }

    // Setup Exception handling
    if (setjmp(png_jmpbuf(png_ptr))) {
	    fprintf(stderr, "Error during png creation\n");
	    code = 1;
	    goto finalise;
    }

    png_init_io(png_ptr, fp);

    // Write header (8 bit colour depth)
    png_set_IHDR(png_ptr, info_ptr, Size, Size,
		    8, PNG_COLOR_TYPE_RGB, PNG_INTERLACE_NONE,
		    PNG_COMPRESSION_TYPE_BASE, PNG_FILTER_TYPE_BASE);

    // Set title
    if (title != NULL) {
	    png_text title_text;
	    title_text.compression = PNG_TEXT_COMPRESSION_NONE;
	    title_text.key = "Title";
	    title_text.text = title;
	    png_set_text(png_ptr, info_ptr, &title_text, 1);
    }

    png_write_info(png_ptr, info_ptr);

    // Allocate memory for one row (3 bytes per pixel - RGB)
    row = (png_bytep) malloc(3 * Size * sizeof(png_byte));

    //define maio e menor elemento
    //********************************
    int i;
    for ( i = 0; i < (Size*Size); i++ )
	if ( buffer[i] > maior )
	    maior = buffer[i];
	    else
		if ( buffer[i] < menor )
		    menor = buffer[i];
    //********************************

    float num;
    for (y=0 ; y<Size ; y++) {
	    for (x=0 ; x<Size ; x++) {
		    num = (float)((float)( buffer[y*Size + x] - menor)/(maior - menor) );
		    setRGB(&(row[x*3]), num );
	    }
	    png_write_row(png_ptr, row);
    }

    // End write
    png_write_end(png_ptr, NULL);

    finalise:
    if (fp != NULL) fclose(fp);
    if (info_ptr != NULL) png_free_data(png_ptr, info_ptr, PNG_FREE_ALL, -1);
    if (png_ptr != NULL) png_destroy_write_struct(&png_ptr, (png_infopp)NULL);
    if (row != NULL) free(row);

    return code;
}

inline void setRGB (png_byte *ptr, float val)
{
	int offset;
	double v = (double)(val * 768);
	if (v > 768) v = 768;

	if ( v != 0 )
	{
	    //printf("Valor Diferente de Zero: %f\n", v );
	    int valor = (int) v;

	    if (valor < 1) valor = 1;
	    offset = valor % 256;

	    if (valor < 256) {
		    ptr[0] = 0; ptr[1] = 0; ptr[2] = offset;
	    }
	    else if (valor < 512) {
		    ptr[0] = 0; ptr[1] = offset; ptr[2] = 255-offset;
	    }
	    else {
		    ptr[0] = offset; ptr[1] = 255-offset; ptr[2] = 0;
	    }
	}
	else
	{
	    v = 0;

	    if ( v == 0 ){
		    ptr[0] = 255; ptr[1] = 255; ptr[2] = 255;
		    //cont++;
	    }
	}
}
